from basicsr.archs import build_network
from basicsr.losses import build_loss
from basicsr.models.sr_model import SRModel
from basicsr.utils import get_root_logger
from basicsr.utils.registry import MODEL_REGISTRY
from collections import OrderedDict
import torch


@MODEL_REGISTRY.register()
class VDSRModel(SRModel):
    """VDSR模型实现，严格按照原始论文设计
    
    原始VDSR使用20个卷积层和残差学习的超分辨率方法。
    训练过程严格遵循论文中的L2损失和优化策略。
    
    Paper: Accurate Image Super-Resolution Using Very Deep Convolutional Networks
    """

    def init_training_settings(self):
        """初始化VDSR的训练设置，严格按照原论文实现，同时保留EMA可选功能"""
        self.net_g.train()
        train_opt = self.opt['train']
        
        # 添加EMA可选分支，符合原始论文的同时提供现代训练选项
        self.ema_decay = train_opt.get('ema_decay', 0)
        if self.ema_decay > 0:
            logger = get_root_logger()
            logger.info(f'使用指数移动平均(EMA)，衰减率: {self.ema_decay}')
            # 定义带EMA的网络，仅用于单GPU测试和保存
            self.net_g_ema = build_network(self.opt['network_g']).to(self.device)
            # 加载预训练模型
            load_path = self.opt['path'].get('pretrain_network_g', None)
            if load_path is not None:
                self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema')
            else:
                self.model_ema(0)  # 复制net_g的权重
            self.net_g_ema.eval()
        
        # 原始VDSR严格使用L2损失(MSE)训练
        if train_opt.get('pixel_opt'):
            self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device)
        else:
            # 确保使用L2损失作为默认值
            l2_opt = {'type': 'MSELoss'}
            self.cri_pix = build_loss(l2_opt).to(self.device)
            logger = get_root_logger()
            logger.warning('未指定像素损失，VDSR使用L2损失作为默认值')
        
        # 原论文不使用感知损失，移除这个选项
        self.cri_perceptual = None
        
        # 设置优化器和学习率调度器
        self.setup_optimizers()
        self.setup_schedulers()

    def optimize_parameters(self, current_iter):
        """执行单步优化，实现原始论文中的训练流程"""
        # 梯度清零
        self.optimizer_g.zero_grad()
        
        # 前向传播
        self.output = self.net_g(self.lq)
        
        # 计算L2损失
        l_total = 0
        loss_dict = OrderedDict()
        
        # 严格使用L2损失
        if self.cri_pix:
            l_pix = self.cri_pix(self.output, self.gt)
            l_total = l_pix  # 原始论文中只有L2损失
            loss_dict['l_pix'] = l_pix
        
        # 反向传播
        l_total.backward()
        
        # 参数更新
        self.optimizer_g.step()
        
        # 记录损失
        self.log_dict = self.reduce_loss_dict(loss_dict)
        
        # EMA模型更新（如果启用）
        if self.ema_decay > 0:
            self.model_ema(decay=self.ema_decay)